package com.hlx.controller;

import com.hlx.entiry.vo.ResultVO;
import com.hlx.enumeration.TypeEnum;
import com.hlx.service.ChatHistoryService;
import com.hlx.service.ChatPdfService;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.vectorstore.QuestionAnswerAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.document.Document;
import org.springframework.ai.reader.ExtractedTextFormatter;
import org.springframework.ai.reader.pdf.PagePdfDocumentReader;
import org.springframework.ai.reader.pdf.config.PdfDocumentReaderConfig;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.core.io.Resource;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;
import org.yaml.snakeyaml.util.UriEncoder;
import reactor.core.publisher.Flux;

import java.io.IOException;
import java.net.URLEncoder;
import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Objects;

@Slf4j
@RequiredArgsConstructor
@RestController
@RequestMapping("/ai/pdf")
public class PdfController {
	private final ChatPdfService chatPdfService;
	private final VectorStore vectorStore;
	private final ChatClient pdfChatClient;
	private final ChatHistoryService chatHistoryService;

	/**
	 * PDF聊天
	 * @param prompt
	 * @param chatId
	 * @return
	 */
	@RequestMapping(value = "/chat", produces = "text/html; chatset=utf-8")
	public Flux<String> chat(String prompt, String chatId){
		chatHistoryService.save(TypeEnum.PDF.getValue(), chatId);
		Resource file = chatPdfService.getFile(chatId);
		return pdfChatClient
				.prompt(prompt)
				.advisors(a -> a.param(ChatMemory.CONVERSATION_ID, chatId))
				.stream()
				.content();
	}

	/**
	 * 上传PDF文件
	 * @param file
	 * @return
	 */
	@PostMapping("/upload")
	public ResultVO uploadPdf(@RequestParam("file") MultipartFile file){
		try {
			// 1.校验文件是否为PDF格式
			if (!file.getContentType().equals("application/pdf")){
				return ResultVO.fail("请上传PDF文件");
			}
			// 2.保存PDF文件
			String chatId = chatPdfService.save(file);
			if (chatId == null || chatId.isEmpty()) {
				return ResultVO.fail("保存PDF文件失败");
			}
			// 3.写入向量库
			wirteToVectorStore(file.getResource());
			return ResultVO.ok(chatId);
		} catch (Exception e){
			log.error("Failed to upload PDF", e);
			return ResultVO.fail("上传PDF文件失败");
		}
	}

	/**
	 * 获取PDF文件
	 * @param chatId
	 * @return
	 * @throws IOException
	 */
	@GetMapping("/file/{chatId}")
	public ResponseEntity<Resource> getFile(@PathVariable("chatId") String chatId) throws IOException {
		// 1.读取文件
		Resource file = chatPdfService.getFile(chatId);
		if (!file.exists()){
			return ResponseEntity.notFound().build();
		}
		// 2.文件名编码，写入响应头
		String fileName = URLEncoder.encode(Objects.requireNonNull(file.getFilename()), StandardCharsets.UTF_8);
		// 3.返回文件
		return ResponseEntity.ok()
				.header("Content-Disposition", "attachment; filename=\"" + fileName + "\"")
				.body(file);
	}

	/**
	 * 将PDF文档写入向量库
	 * @param resource
	 */
	private void wirteToVectorStore(Resource resource) {
		// 创建PDF文档读取器
		PagePdfDocumentReader pdfReader = new PagePdfDocumentReader(resource,
				PdfDocumentReaderConfig.builder()
						.withPageExtractedTextFormatter(ExtractedTextFormatter.defaults())
						.withPagesPerDocument(1)
						.build());
		// 读取PDF文档
		List<Document> documents = pdfReader.read();
		// 将文档写入向量库
		vectorStore.add(documents);
	}
}
